#!/usr/bin/env node
/**
 * Node.js WebSocket测试客户端
 * 用于测试AI聊天流式输出的WebSocket连接
 */

import WebSocket from 'ws';
import { createInterface } from 'readline';
import { fileURLToPath } from 'url';
import { dirname } from 'path';

class WebSocketTestClient {
    constructor(baseUrl = 'ws://localhost:8000', conversationId = 1) {
        this.baseUrl = baseUrl;
        this.conversationId = conversationId;
        this.wsUrl = `${baseUrl}/ws/ai-chat/conversations/${conversationId}/stream/`;
        this.ws = null;
        this.isConnected = false;
    }

    connect() {
        return new Promise((resolve, reject) => {
            console.log(`🔗 尝试连接到: ${this.wsUrl}`);
            
            try {
                this.ws = new WebSocket(this.wsUrl);
                
                this.ws.on('open', () => {
                    console.log('✅ WebSocket连接成功建立！');
                    this.isConnected = true;
                    this.setupEventHandlers();
                    resolve(true);
                });
                
                this.ws.on('error', (error) => {
                    console.error(`❌ WebSocket连接失败:`, error.message);
                    this.isConnected = false;
                    reject(error);
                });
                
            } catch (error) {
                console.error(`❌ WebSocket连接异常:`, error.message);
                reject(error);
            }
        });
    }

    setupEventHandlers() {
        if (!this.ws) return;

        this.ws.on('message', (data) => {
            try {
                const message = JSON.parse(data.toString());
                this.handleMessage(message);
            } catch (error) {
                console.error('❌ 解析JSON失败:', data.toString());
            }
        });

        this.ws.on('close', (code, reason) => {
            console.log(`🔌 WebSocket连接已关闭 - 代码: ${code}, 原因: ${reason || '无'}`);
            this.isConnected = false;
        });

        this.ws.on('error', (error) => {
            console.error('❌ WebSocket错误:', error.message);
            this.isConnected = false;
        });
    }

    handleMessage(data) {
        const messageType = data.type;
        const timestamp = new Date().toLocaleTimeString();
        
        console.log(`📥 [${timestamp}] 收到消息类型: ${messageType}`);
        
        switch (messageType) {
            case 'stream_start':
                console.log(`🚀 流式对话开始 - Stream ID: ${data.stream_id}`);
                break;
                
            case 'stream_chunk':
                const content = data.content || '';
                const tokens = data.tokens || 0;
                const isFinal = data.is_final || false;
                const displayContent = content.length > 50 ? content.substring(0, 50) + '...' : content;
                console.log(`📝 流式内容: ${displayContent}`);
                console.log(`   Token数: ${tokens}, 是否最终: ${isFinal}`);
                break;
                
            case 'stream_complete':
                console.log(`✅ 流式对话完成 - 消息ID: ${data.message_id}, 总Token: ${data.total_tokens}`);
                break;
                
            case 'stream_error':
                console.log(`❌ 流式对话错误: ${data.message}`);
                break;
                
            case 'stream_stopped':
                console.log(`🛑 流式对话已停止 - Stream ID: ${data.stream_id}`);
                break;
                
            case 'error':
                console.log(`❌ 服务器错误: ${data.message}`);
                break;
                
            case 'pong':
                console.log('🏓 收到心跳响应');
                break;
                
            default:
                console.log('❓ 未知消息类型:', data);
        }
    }

    sendMessage(messageData) {
        if (!this.isConnected || !this.ws) {
            console.error('❌ WebSocket未连接');
            return false;
        }

        try {
            const messageJson = JSON.stringify(messageData);
            this.ws.send(messageJson);
            console.log(`📤 发送消息: ${messageJson}`);
            return true;
        } catch (error) {
            console.error('❌ 发送消息失败:', error.message);
            return false;
        }
    }

    sendPing() {
        const pingData = {
            type: 'ping',
            timestamp: Date.now()
        };
        return this.sendMessage(pingData);
    }

    startStream(message, model = 'gpt-3.5-turbo', role = 'user') {
        const streamData = {
            type: 'start_stream',
            message: message,
            model: model,
            role: role
        };
        return this.sendMessage(streamData);
    }

    stopStream(streamId) {
        const stopData = {
            type: 'stop_stream',
            stream_id: streamId
        };
        return this.sendMessage(stopData);
    }

    disconnect() {
        if (this.ws) {
            this.ws.close();
            console.log('🔌 WebSocket连接已断开');
        }
        this.isConnected = false;
    }
}

// 自动测试函数
async function autoTest() {
    console.log('🧪 开始WebSocket自动连接测试');
    console.log('='.repeat(50));
    
    const client = new WebSocketTestClient();
    
    try {
        // 1. 建立连接
        await client.connect();
        
        // 2. 等待连接稳定
        await new Promise(resolve => setTimeout(resolve, 1000));
        
        // 3. 发送心跳测试
        console.log('\n🏓 测试心跳检测...');
        client.sendPing();
        await new Promise(resolve => setTimeout(resolve, 2000));
        
        // 4. 测试流式对话
        console.log('\n🚀 测试流式对话...');
        client.startStream('你好，请简单介绍一下你自己');
        
        // 5. 等待响应
        console.log('\n⏳ 等待响应（30秒）...');
        await new Promise(resolve => setTimeout(resolve, 30000));
        
        // 6. 再次测试心跳
        console.log('\n🏓 再次测试心跳...');
        client.sendPing();
        await new Promise(resolve => setTimeout(resolve, 2000));
        
    } catch (error) {
        console.error('❌ 测试过程中出错:', error.message);
    } finally {
        client.disconnect();
        console.log('🧹 测试完成');
    }
}

// 交互式测试函数
async function interactiveTest() {
    console.log('🎮 WebSocket交互式测试模式');
    console.log('='.repeat(50));
    
    const client = new WebSocketTestClient();
    
    try {
        await client.connect();
        
        // 等待连接稳定
        await new Promise(resolve => setTimeout(resolve, 1000));
        
        console.log('\n可用命令:');
        console.log('  ping - 发送心跳');
        console.log('  stream <消息> - 开始流式对话');
        console.log('  stop <stream_id> - 停止流式对话');
        console.log('  quit - 退出');
        console.log();
        
        const rl = createInterface({
            input: process.stdin,
            output: process.stdout
        });
        
        const promptUser = () => {
            rl.question('👤 输入命令: ', (command) => {
                const trimmedCommand = command.trim();
                
                if (trimmedCommand.toLowerCase() === 'quit') {
                    rl.close();
                    client.disconnect();
                    process.exit(0);
                } else if (trimmedCommand.toLowerCase() === 'ping') {
                    client.sendPing();
                    promptUser();
                } else if (trimmedCommand.startsWith('stream ')) {
                    const message = trimmedCommand.substring(7);
                    client.startStream(message);
                    promptUser();
                } else if (trimmedCommand.startsWith('stop ')) {
                    const streamId = trimmedCommand.substring(5);
                    client.stopStream(streamId);
                    promptUser();
                } else {
                    console.log('❓ 未知命令，请重试');
                    promptUser();
                }
            });
        };
        
        promptUser();
        
    } catch (error) {
        console.error('❌ 交互测试出错:', error.message);
    }
}

// 主函数
function main() {
    console.log('Node.js WebSocket测试客户端');
    console.log('请选择测试模式:');
    console.log('1. 自动测试');
    console.log('2. 交互式测试');
    
    const rl = createInterface({
        input: process.stdin,
        output: process.stdout
    });
    
    rl.question('请输入选择 (1/2): ', (choice) => {
        rl.close();
        
        if (choice === '1') {
            autoTest().catch(console.error);
        } else if (choice === '2') {
            interactiveTest().catch(console.error);
        } else {
            console.log('❌ 无效选择');
        }
    });
}

// 处理程序退出
process.on('SIGINT', () => {
    console.log('\n👋 测试已取消');
    process.exit(0);
});

process.on('uncaughtException', (error) => {
    console.error('❌ 程序出错:', error.message);
    process.exit(1);
});

// ES模块中检查是否为主模块
const __filename = fileURLToPath(import.meta.url);
const __dirname = dirname(__filename);

// 简化的主模块检查
const isMain = process.argv[1] && process.argv[1].endsWith('test_websocket_client.js');
if (isMain) {
    main();
}

export default WebSocketTestClient;